# rsa 加密 解密
import base64

import rsa

from zz import zz_log


def file_to_public_key(file_path: str) -> rsa.PublicKey:
    """
    文件转公钥对象
    :param file_path: 文件路径，test.pem
    :return:
    """
    with open(file_path, "rb") as file:
        data = file.read()
        key = rsa.PublicKey.load_pkcs1(data)
    return key


def file_to_private_key(file_path: str) -> rsa.PrivateKey:
    """
    文件转私钥对象
    :param file_path: 文件路径，test.pem
    :return:
    """
    with open(file_path, "rb") as file:
        data = file.read()
        key = rsa.PrivateKey.load_pkcs1(data)
    return key


def str_to_public_key(s: str) -> rsa.PublicKey:
    """
    字符串转公钥对象
    :param s:
    :return:
    """
    return rsa.PublicKey.load_pkcs1(s.encode())


def str_to_private_key(s: str) -> rsa.PrivateKey:
    """
    字符串转私钥对象
    :param s:
    :return:
    """
    return rsa.PrivateKey.load_pkcs1(s.encode())


def encode(data, public_key: rsa.PublicKey) -> str:
    """
    加密返回字符串
    :param data:
    :param public_key:
    :return:
    """
    if isinstance(data, str):
        data = data.encode('utf-8')
    encrypt_data = rsa.encrypt(data, public_key)
    return base64.b64encode(encrypt_data).decode()


def decode(data, private_key: rsa.PrivateKey) -> str:
    """
    解密返回字符串
    :param data:
    :param private_key:
    :return:
    """
    if isinstance(data, str):
        data = base64.b64decode(data)
    return rsa.decrypt(data, private_key).decode()


def sign(data, private_key: rsa.PrivateKey) -> str:
    """
    生成签名
    :param data:
    :param private_key:
    :return:
    """
    if isinstance(data, str):
        data = data.encode('utf-8')
    sign_data = rsa.sign(data, private_key, "SHA-256")
    return base64.b64encode(sign_data).decode()


def sign_check(data, sign, public_key: rsa.PublicKey) -> bool:
    """
    检查签名
    :param data:
    :param sign:
    :param public_key:
    :return:
    """
    if isinstance(data, str):
        data = data.encode('utf-8')
    if isinstance(sign, str):
        sign = base64.b64decode(sign)

    try:
        v = rsa.verify(data, sign, public_key)
        return v == "SHA-256"
    except Exception as ex:
        zz_log.of().error(ex)
    return False
